package rerac.feature;

import java.util.Map;

import rerac.protos.Corpus.Document;
import rerac.protos.Corpus.Document.Annotation;
import rerac.protos.Corpus.Document.Compound;
import rerac.protos.Corpus.Document.CompoundGroup;
import rerac.protos.Corpus.Document.Method;
import rerac.protos.Corpus.Document.Token;
import rerac.protos.Corpus.Document.AnnotationType;
import rerac.protos.Corpus.Document.Compound.CompoundSlot;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;

/**
 * This adds, for single-sentence documents, the information whether a pair of
 * tokens, that are known arguments for certain relations, matches or not.
 * 
 * @author Benjamin Roth
 *
 */
public class SentenceExtractor {
  
  Map<String, Multimap<String, String>> rel2seeds;

  /**
   * This constructs a extractor that marks annotation matches for argument
   * pairs of a given relation.
   * 
   * The argument pairs (seed tokens) are grouped by relation, and passed by
   * a map from relations to pairs. The pairs themselves are contained in a 
   * MultiMap from argument_1 to argument_2.
   * 
   * @param seedTokens, the Map from relation to token pairs.
   */
  public SentenceExtractor(Map<String, Multimap<String, String>> seedTokens) {
    rel2seeds = seedTokens;
  }
  
  /**
   * This adds an annotation to the document if a pair of seed token matches.
   * A CompoundGroup of type AnnotationType.PROPERTY is added, and compound 
   * for every match are added. The text of the Compound is the name of the 
   * relation, and the first and second slot of it point to the indices of
   * the first and second argument of the relation respectively.
   * 
   * @param doc the document for which matches should be annotated.
   * @return the annotated document.
   */
  public Document extract(Document doc) {
    Document.Builder extractedDoc = Document.newBuilder(doc);
    
    int methodIdx = doc.getMethodCount();
    AnnotationType methodType = AnnotationType.PROPERTY;
    Method.Builder m = Method.newBuilder();
    m.setId(SentenceExtractor.class.getName());
    m.setType(methodType);
    extractedDoc.addMethod(m);
    
    int compoundGroupIdx = doc.getCompoundCount();
    CompoundGroup.Builder extractions = CompoundGroup.newBuilder();
    extractions.setMethodIndex(methodIdx);
    extractions.setType(methodType);
    
    Multimap<String, Integer> token2idx = HashMultimap.create();
    for (int i = 0; i < doc.getTokenCount(); ++i) {
      token2idx.put(doc.getToken(i).getText(), new Integer(i));
    }
    for (int tokIdx = 0; tokIdx < doc.getTokenCount(); ++tokIdx) {
      for (String relation : rel2seeds.keySet()) {
        Multimap<String, String> arg2arg = rel2seeds.get(relation);
        String tok = doc.getToken(tokIdx).getText();
        for (String target : arg2arg.get(tok)) {
          if (token2idx.containsKey(target)) {
            for (int targetIdx : token2idx.get(target)){
              Compound.Builder relationCompound = Compound.newBuilder();
              relationCompound.addSlot(
                  CompoundSlot.newBuilder().setTokenIndex(tokIdx));
              relationCompound.addSlot(
                  CompoundSlot.newBuilder().setTokenIndex(targetIdx));
              relationCompound.setText(relation);
              int compoundIdx = extractions.getCompoundCount();
              extractions.addCompound(relationCompound);
              Token.Builder tokB = Token.newBuilder(extractedDoc.getToken(tokIdx));
              tokB.addAnnotation(Annotation.newBuilder()
                  .setMethodIndex(methodIdx)
                  .setType(methodType)
                  .setCompoundGroup(compoundGroupIdx)
                  .setTargetCompound(compoundIdx)
                  .build());
              extractedDoc.setToken(tokIdx, tokB);
              Token.Builder targetB = Token.newBuilder(extractedDoc.getToken(targetIdx));
              targetB.addAnnotation(Annotation.newBuilder()
                  .setMethodIndex(methodIdx)
                  .setType(methodType)
                  .setCompoundGroup(compoundGroupIdx)
                  .setTargetCompound(compoundIdx)
                  .build());
              extractedDoc.setToken(targetIdx, targetB);
            }
          }
        }
      }
    }
    extractedDoc.addCompound(extractions);
    return extractedDoc.build();
  }
  
  /**
   * This returns true if the document has annotations by this class.
   * 
   * @param doc the document to check.
   * @return true if there is an annotation for matching relation arguments.
   */
  public static boolean hasAnnotation(Document doc) {
    boolean hasAnno = false;
    int methodIdx = -1;
    for (int i = 0; i < doc.getMethodCount() && methodIdx == -1; ++i) {
      Method m = doc.getMethod(i);
      if (m.hasId() && m.getId().equals(SentenceExtractor.class.getName())) {
        methodIdx = i;
      }
    }
    for (CompoundGroup cg : doc.getCompoundList()) {
      if (cg.getMethodIndex() == methodIdx) {
        hasAnno = (cg.getCompoundCount() > 0);
        break;
      }
    }
    return hasAnno;
  }

}
